# 通用的fit函数，输入是dataloader，进行了模型训练，并且返回了loss和acc
import torch
from Config import config   #从自定义的Config.py参数文件中插入

# 加载参数
# 导入之前的Config.py文件中的config类，进行实例化，产生【对象(因为只有对象才能调用类中的属性和方法)】
config = config()
# 自定义fit函数，该函数接受的参数为：
#   epoch 当前的训练轮次
#   model 要训练的深度学习模型
#   trainloader 训练数据集的DataLoader 对象，用来迭代获取批次的数据
#   loss_fn 损失函数，用于计算预测值与真实值之间的误差
#   optim 优化器，用于更新模型参数以最小化损失函数
#   device 设备信息(如cpu或cuda)，指定模型和数据应该在哪个设备上运行
#   bst_loss 最佳损失值，用于保存最佳模型
def fit(epoch, model, trainloader, testloader, loss_fn, optim, device, bst_loss):
    correct = 0     #正确预测的数量
    total = 0       #总样本数量
    running_loss = 0    #运行中的累计损失，用于计算平均损失
    for x, y in trainloader:    #遍历trainloader中的每一个batch，每个batch包含输入数据x和对应的标签y
        x, y = x.to(device), y.to(device)   #将输入数据x和标签y移动到指定的设备(GPU或CPU)
        y_pred = model(x)   #将输入数据x传递给模型model,得到预测输出y_pred,这是向前传播的步骤
        loss = loss_fn(y_pred, y)   #使用损失函数loss_fn计算预测输出y_pred和真实标签y之间的损失loss
        optim.zero_grad()      #清空之前的梯度，因为pytorch会默认累加梯度，所以在每次更新参数之前需要将梯度清零
        loss.backward()     #执行反向传播，计算损失值对于每个参数的梯度
        optim.step()    #根据计算出的梯度更新模型参数，这是由优化器optim完成的
        with torch.no_grad():   #进入一个上下文管理器，在当前上下文中，所有的张量操作都不会记录梯度信息，这是为了节省内存
            # torch.orgmax函数返回y_pred在指定维度上的最大值的索引，即预测类别标签
            y_pred = torch.argmax(y_pred, dim=1)    #对于分类任务，通常需要找到每个样本预测概率最 的类别
            # y_pred == y 生成一个布尔张量(True和False),表示哪些是正确的
            # .sum() 计算其中True的数量
            # .item()将结果转换为python标量，每次的标量都会累加存储到correct容器当中
            correct += (y_pred == y).sum().item()   #统计正确的样本数量
            # total变量，用于记录处理过的总样本数
            # y.size(0)返回当前batch的大小，也就是该batch中样本的数量
            total += y.size(0)  #样本总数
            # loss.item() 提取损失张量中的标量值
            running_loss += loss.item()
# #---------------------------------------------------------------------------------#
# #--------------epoch_acc 和 epoch_loss的计算放在循环外，表示他们只会在最后一个batch计算----#
# #---------------------------------------------------------------------------------#
# 计算当前epoch的准确率epoch_acc,即正确预测的样本数除以总样本数
    epoch_acc = correct / total
# 计算当前epoch的平均损失epoch_loss, 即将累计损失除以测试数据集的大小
# 注意len(trainloader.dataset)指的是整个测试数据集的大小，而不是batch数量，如果是batch数量，则应该除以len(trainloader)
    epoch_loss = running_loss / len(trainloader.dataset)

# #---------------------------------------------------------------------------------#
# #-----------用于评估在测试集上的性能代码模块，它与训练过程中的评估类似，但不涉及参数更新---------#
# #---------------------------------------------------------------------------------#
    test_correct = 0        #累计正确预测的样本数量
    test_total = 0          #累计测试数据集中所有样本的数量
    test_running_loss = 0   #累计测试过程中的所有批次的损失值
    with torch.no_grad():
        for x, y  in testloader:
            x, y = x.to(device), y.to(device)
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)  # 将输入数据x传递给模型model,得到预测输出y_pred,这是向前传播的步骤
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()
    test_epoch_acc = test_correct / test_total
    test_epoch_loss = test_running_loss / len(testloader.dataset)

    if test_epoch_loss < bst_loss:      #保存验证集最优模型
        bst_loss = test_epoch_loss
        torch.save(model.state_dict(), config.save_path)
    print('epoch:', epoch, ', loss:', round(epoch_loss, 5),
          ', accuracy:', round(epoch_acc, 5),
          '，test_loss:', round(test_epoch_loss, 5),
          ', test_accuracy:', round(test_epoch_acc,5)
          )
    return epoch_loss, epoch_acc, test_epoch_loss, test_epoch_acc